import os
import math
import argparse
import random
import logging
import options.options as option
from utils import util
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
from data.data_sampler import DistIterSampler
from data import create_dataloader, create_dataset
from models import create_model
from math import ceil


def init_dist(backend='nccl', **kwargs):
    ''' initialization for distributed training'''
    # if mp.get_start_method(allow_none=True) is None:
    if mp.get_start_method(allow_none=True) != 'spawn':
        mp.set_start_method('spawn')
    rank = int(os.environ['RANK'])
    num_gpus = torch.cuda.device_count()
    torch.cuda.set_device(rank % num_gpus)
    dist.init_process_group(backend=backend, **kwargs)


def main():
    #### options
    parser = argparse.ArgumentParser()
    parser.add_argument('-opt', type=str, default='options/train/train_InvDN_aug_mcj.yml',help='Path to option YMAL file.')
    parser.add_argument('--launcher', choices=['none', 'pytorch'], default='none',
                        help='job launcher')
    parser.add_argument('--local_rank', type=int, default=0)
    args = parser.parse_args()
    opt = option.parse(args.opt, is_train=True)

    # test scale x4 or x8
    opt['model'] = 'InvDN_C'
    opt['scale']= 4
    opt['gpu_ids']=[6] #7
    opt['datasets']['train']['GT_size']=32
    opt['datasets']['train']['scale']=4
    opt['datasets']['train']['mode']= 'LQGTRN_Seis_train'
    opt['datasets']['train']['noise_mode'] = 'B'
    opt['datasets']['train']['data_mul'] = ceil(1000/3) # 1000

    opt['datasets']['train']['data_aug'] = True
    opt['datasets']['train']['fake_ratio'] = 2
    opt['datasets']['train']['epsilon'] = 2e-4
    # opt['datasets']['train']['epsilon'] = [1e-3, 2e-3, 5e-3, 1e-2]  # [1e-4, 2e-4, 5e-4, 1e-3]
    opt['datasets']['train']['data_mul_aft_aug']= ceil(opt['datasets']['train']['data_mul']
                                                      * (opt['datasets']['train']['fake_ratio'] + 1))

    opt['datasets']['train']['batch_size'] = 64  # len(dataset)=data_mul*batch_size
    opt['datasets']['train']['sigma'] = 75
    opt['datasets']['train']['dataset_type'] = 'orthofxdeconXJ10'
    opt['datasets']['val']['scale'] = 4
    opt['datasets']['val']['mode'] = 'LQGTRN_Seis_val'
    opt['train']['pixel_criterion_back']='l1' #default l1
    opt['network_G']['scale']= 4
    opt['network_G']['block_num'] = [8,8]
    opt['split_channel']=3
    opt['train']['val_freq']=opt['datasets']['train']['data_mul_aft_aug']
    epoch = 60
    opt['train']['niter'] = opt['datasets']['train']['data_mul_aft_aug'] * epoch
    opt['train']['lr_steps'] = [i * opt['datasets']['train']['data_mul_aft_aug'] for i in [10, 20, 25, 30, 35, 40, 45, 50]]  # milestone
    opt['train']['lambda_SSIM_Grad'] = 0
    opt['logger']['save_checkpoint_freq']=opt['datasets']['train']['data_mul_aft_aug']
    # opt['path']['resume_state']=''

    # opt['path']['model_generator'] = '../experiments/expert_XJ/InvDN_ResUnit_x4_expert_XJ10/models/10_G.pth'
    opt['path']['model_generator'] = '../experiments/orfx_XJ/InvDN_ResUnit_x4_orfx_XJ10_sc3/models/3_G.pth'

    import os.path as osp
    opt['name'] = 'InvDN_ResUnit_x4_XJ10_aug2ep2e4' #InvDN_ResUnit_x4_orfx_XJ30_1 sc3
    experiments_root = osp.join(opt['path']['root'], 'experiments/orfx_XJ', opt['name'])
    opt['path']['experiments_root'] = experiments_root
    opt['path']['models'] = osp.join(experiments_root, 'models')
    opt['path']['training_state'] = osp.join(experiments_root, 'training_state')
    opt['path']['log'] = experiments_root
    opt['path']['val_images'] = osp.join(experiments_root, 'val_images')

    #### distributed training settings
    if args.launcher == 'none':  # disabled distributed training
        opt['dist'] = False
        rank = -1
        print('Disabled distributed training.')
    else:
        opt['dist'] = True
        init_dist()
        world_size = torch.distributed.get_world_size()
        rank = torch.distributed.get_rank()

    #### loading resume state if exists
    if opt['path'].get('resume_state', None):
        # distributed resuming: all load into default GPU
        device_id = torch.cuda.current_device()
        resume_state = torch.load(opt['path']['resume_state'],
                                  map_location=lambda storage, loc: storage.cuda(device_id))
        option.check_resume(opt, resume_state['iter'])  # check resume options
    else:
        resume_state = None

    #### mkdir and loggers
    if rank <= 0:  # normal training (rank -1) OR distributed training (rank 0)
        if resume_state is None:
            util.mkdir_and_rename(
                opt['path']['experiments_root'])  # rename experiment folder if exists
            util.mkdirs((path for key, path in opt['path'].items() if not key == 'experiments_root'
                         and 'pretrain_model' not in key and 'resume' not in key))

        # config loggers. Before it, the log will not work
        util.setup_logger('base', opt['path']['log'], 'train_' + opt['name'], level=logging.INFO,
                          screen=True, tofile=True)
        util.setup_logger('val', opt['path']['log'], 'val_' + opt['name'], level=logging.INFO,
                          screen=True, tofile=True)
        logger = logging.getLogger('base')
        logger.info(option.dict2str(opt))
        # tensorboard logger
        if opt['use_tb_logger'] and 'debug' not in opt['name']:
            version = float(torch.__version__[0:3])
            if version >= 1.1:  # PyTorch 1.1
                from torch.utils.tensorboard import SummaryWriter
            else:
                logger.info(
                    'You are using PyTorch {}. Tensorboard will use [tensorboardX]'.format(version))
                from tensorboardX import SummaryWriter
            tb_logger = SummaryWriter(log_dir='../tb_logger/' + opt['name'])
    else:
        util.setup_logger('base', opt['path']['log'], 'train', level=logging.INFO, screen=True)
        logger = logging.getLogger('base')

    # convert to NoneDict, which returns None for missing keys
    opt = option.dict_to_nonedict(opt)

    #### random seed
    seed = opt['train']['manual_seed']
    if seed is None:
        seed = random.randint(1, 10000)
    if rank <= 0:
        logger.info('Random seed: {}'.format(seed))
    util.set_random_seed(seed)

    torch.backends.cudnn.benchmark = True
    # torch.backends.cudnn.deterministic = True

    #### create train and val dataloader
    dataset_ratio = 200  # enlarge the size of each epoch
    for phase, dataset_opt in opt['datasets'].items():
        if phase == 'train':
            train_set = create_dataset(dataset_opt)
            train_size = int(math.ceil(len(train_set) / dataset_opt['batch_size']))
            total_iters = int(opt['train']['niter'])
            total_epochs = int(math.ceil(total_iters / train_size))
            if opt['dist']:
                train_sampler = DistIterSampler(train_set, world_size, rank, dataset_ratio)
                total_epochs = int(math.ceil(total_iters / (train_size * dataset_ratio)))
            else:
                train_sampler = None
            train_loader = create_dataloader(train_set, dataset_opt, opt, train_sampler)
            if rank <= 0:
                logger.info('Number of train images: {:,d}, iters: {:,d}'.format(
                    len(train_set), train_size))
                logger.info('Total epochs needed: {:d} for iters {:,d}'.format(
                    total_epochs, total_iters))
        elif phase == 'val':
            val_set = create_dataset(dataset_opt)
            val_loader = create_dataloader(val_set, dataset_opt, opt, None)
            if rank <= 0:
                logger.info('Number of val images in [{:s}]: {:d}'.format(
                    dataset_opt['name'], len(val_set)))
        else:
            raise NotImplementedError('Phase [{:s}] is not recognized.'.format(phase))
    assert train_loader is not None

    #### create model
    model = create_model(opt)
    # model_generater= create_model(opt)

    #### resume training
    if resume_state:
        logger.info('Resuming training from epoch: {}, iter: {}.'.format(
            resume_state['epoch'], resume_state['iter']))

        start_epoch = resume_state['epoch']
        current_step = resume_state['iter']
        model.resume_training(resume_state)  # handle optimizers and schedulers
    else:
        current_step = 0
        start_epoch = 0

    #### training
    logger.info('Start training from epoch: {:d}, iter: {:d}'.format(start_epoch, current_step))
    for epoch in range(start_epoch, total_epochs + 1):
        if opt['dist']:
            train_sampler.set_epoch(epoch)
        for _, train_data in enumerate(train_loader):
            current_step += 1
            if current_step > total_iters:
                break
            # ### augment dataset
            # with torch.autograd.no_grad():
            #     im_noisy[NReal:, ] = sample_generator(netG, im_gt[NReal:, ])
            #     im_noisy[NReal:, ].clamp_(-1.0, 1.0)
            # optimizer.zero_grad()

            #### training
            # model.feed_data(train_data)
            model.feed_data_aug(train_data)#,model_generater,epsilon
            model.optimize_parameters(current_step)

            #### update learning rate
            model.update_learning_rate(current_step, warmup_iter=opt['train']['warmup_iter'])

            #### log
            if current_step % opt['logger']['print_freq'] == 0:
                logs = model.get_current_log()
                message = '<epoch:{:3d}, iter:{:8,d}, lr:{:.3e}> '.format(
                    epoch, current_step, model.get_current_learning_rate())
                for k, v in logs.items():
                    message += '{:s}: {:.4e} '.format(k, v)
                    # tensorboard logger
                    if opt['use_tb_logger'] and 'debug' not in opt['name']:
                        if rank <= 0:
                            tb_logger.add_scalar(k, v, current_step)
                if rank <= 0:
                    logger.info(message)

            # validation
            if current_step % opt['train']['val_freq'] == 0 and rank <= 0:
                avg_psnr = 0.0
                avg_ssim =0.0
                idx = 0
                for val_data in val_loader:
                    idx += 1
                    # img_name = os.path.splitext(os.path.basename(val_data['LQ_path'][0]))[0] 如果是切好片的
                    img_name= val_data['img_index'][0]
                    img_dir = os.path.join(opt['path']['val_images'], img_name)
                    util.mkdir(img_dir)

                    model.feed_data(val_data)
                    model.test()

                    visuals = model.get_current_visuals()

                    # calculate PSNR
                    crop_size = opt['scale']
                    gt= visuals['GT'].clamp_(-1.0, 1.0).numpy()
                    denoised = visuals['Denoised'].clamp_(-1.0, 1.0).numpy()
                    cropped_denoised = denoised[:, crop_size:-crop_size, crop_size:-crop_size]
                    cropped_gt = gt[ :, crop_size:-crop_size, crop_size:-crop_size]
                    from skimage.metrics import peak_signal_noise_ratio as compare_psnr, structural_similarity as compare_ssim
                    avg_psnr += compare_psnr(cropped_denoised.squeeze(), cropped_gt.squeeze())
                    avg_ssim += compare_ssim(cropped_denoised.squeeze(), cropped_gt.squeeze())


                    # # save image
                    # denoised_img = util.tensor2img(visuals['Denoised'])  # uint8
                    # gt_img = util.tensor2img(visuals['GT'])  # uint8
                    # noisy_img = util.tensor2img(visuals['Noisy']) # uint8
                    #
                    # lr_img = util.tensor2img(visuals['LR'])
                    # gtl_img = util.tensor2img(visuals['LR_ref'])

                    # if current_step % (opt['train']['val_freq']) == 0: #每10个epoch保存一次 *10
                    #     # Save Denoised images
                    #     save_img_path = os.path.join(img_dir,
                    #                                  '{:s}_{:d}.png'.format(img_name, current_step))
                    #     util.save_img(denoised_img, save_img_path)
                    #
                    #     # Save LR images
                    #     save_img_path_L = os.path.join(img_dir, '{:s}_forwLR_{:d}.png'.format(img_name, current_step))
                    #     util.save_img(lr_img, save_img_path_L)
                    #
                    # # Save ground truth
                    # if current_step == opt['train']['val_freq']:
                    #     save_img_path_gt = os.path.join(img_dir, '{:s}_GT_{:d}.png'.format(img_name, current_step))
                    #     util.save_img(gt_img, save_img_path_gt)
                    #     save_img_path_noisy = os.path.join(img_dir, '{:s}_Noisy_{:d}.png'.format(img_name, current_step))
                    #     util.save_img(noisy_img, save_img_path_noisy)
                    #     save_img_path_gtl = os.path.join(img_dir, '{:s}_LR_ref_{:d}.png'.format(img_name, current_step))
                    #     util.save_img(gtl_img, save_img_path_gtl)

                avg_psnr = avg_psnr / idx
                avg_ssim = avg_ssim / idx

                # log
                logger.info('# Validation # PSNR={:4.2f}, SSIM: {:5.4f}.'.format(avg_psnr,avg_ssim))
                logger_val = logging.getLogger('val')  # validation logger
                logger_val.info('<epoch:{:3d}, iter:{:8,d}> psnr: {:.4e}, SSIM: {:5.4f}.'.format(
                    epoch, current_step, avg_psnr, avg_ssim))
                # tensorboard logger
                if opt['use_tb_logger'] and 'debug' not in opt['name']:
                    tb_logger.add_scalar('psnr', avg_psnr, current_step)
                    tb_logger.add_scalar('ssim', avg_ssim, current_step)

            #### save models and training states
            if current_step % opt['logger']['save_checkpoint_freq'] == 0:
                if rank <= 0:
                    logger.info('Saving models and training states.')
                    model.save(int(current_step//opt['logger']['save_checkpoint_freq']))
                    model.save_training_state(epoch, int(current_step//opt['logger']['save_checkpoint_freq']))

    if rank <= 0:
        logger.info('Saving the final model.')
        model.save('latest')
        logger.info('End of training.')


if __name__ == '__main__':
    main()
